package main

import (
	"3nets.io/wireguardcommon"
	"crypto"
	"encoding/base64"
	"fmt"
	"github.com/vishvananda/netlink"
	"net"
	"net/netip"
	"os"
	"strings"
)

const (
	ConnHost = "0.0.0.0"
	ConnPort = "32323"
	ConnType = "tcp"
)

var managerPrivate crypto.PrivateKey
var managerPublic crypto.PublicKey
var wireguardLink netlink.Link

func main() {
	var err error

	managerPrivate, managerPublic, _, wireguardLink, err = wireguardcommon.MakeWireguardInterface()

	if err != nil {
		fmt.Println("Could not make Wireguard interface, err:", err.Error())
		os.Exit(1)
	}

	// Listen for incoming connections.
	l, err := net.Listen(ConnType, ConnHost+":"+ConnPort)
	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer func(l net.Listener) {
		err := l.Close()
		if err != nil {
			fmt.Println("error closing listener on " + ConnHost + ":" + ConnPort)
		}
	}(l)
	fmt.Println("Listening on " + ConnHost + ":" + ConnPort)
	for {
		// Listen for an incoming connection.
		conn, err := l.Accept()
		if err != nil {
			fmt.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		go handleRequest(conn)
	}
}

// Handles incoming requests.
func handleRequest(conn net.Conn) {
	maxReqLen := 256
	edgeResponse, responseLen, err := wireguardcommon.ReadConn(conn, maxReqLen)
	if err != nil {
		return
		fmt.Printf("ReadConn retured error : %s\n", edgeResponse)
	}

	fmt.Printf("edgeResponse: %s\n", edgeResponse)
	edgeSignature := edgeResponse[:responseLen]
	instanceId := ""
	signatureTokens := strings.Split(edgeResponse, ":")
	if len(signatureTokens) > 1 {
		edgeSignature = signatureTokens[0]
		instanceId = signatureTokens[1]
		fmt.Printf("edgeSignature: %s\n", edgeSignature)
		fmt.Printf("instanceId: %s\n", instanceId)
		if !wireguardcommon.VerifyVmInstanceId(instanceId) {
			wireguardcommon.CloseConn(conn)
			return
		}
	}

	//check against "3netsedge.io"
	_, ok, err := wireguardcommon.VerifySignature("", edgeSignature)
	if !ok || err != nil {

		fmt.Printf("Signature didnot match : 3netsedge.io\n")
		//check against "*3netsedge.io"
		ok, err := wireguardcommon.VerifySignatureRegex("", ".*3netsedge.io", edgeSignature)
		fmt.Printf("ok = %s, err= %s\n", ok, err)
		if !ok || err != nil {
			fmt.Printf("Signature didnot match : *3netsedge.io\n")
			fmt.Printf("Edge Signature not ok error : %s\n", err)
			wireguardcommon.CloseConn(conn)
			return
		}
	}

	var managerPrivateBytes [32]byte
	var managerPublicBytes [32]byte
	if ok := wireguardcommon.CheckType(&managerPrivateBytes, managerPrivate); !ok {
		panic("ecdh: unexpected type of manager private key")
	}
	if ok := wireguardcommon.CheckType(&managerPublicBytes, managerPublic); !ok {
		panic("ecdh: unexpected type of manager public key")
	}

	managerDHPrivate := base64.StdEncoding.EncodeToString(managerPrivateBytes[:])
	managerDHPublic := base64.StdEncoding.EncodeToString(managerPublicBytes[:])
	fmt.Printf("managerDHPublic: %s\n", managerDHPublic)
	fmt.Printf("managerDHPrivate: %s\n", managerDHPrivate)

	// Send manager DH public
	writeLen, err := wireguardcommon.WriteConn(conn, []byte(managerDHPublic))
	if err != nil || writeLen != len(managerDHPublic) {
		fmt.Printf("write length mismatch. expected = %d, got=%d\n", len(managerDHPublic), writeLen)
		return
	}

	edgeDHPublic, responseLen, err := wireguardcommon.ReadConn(conn, maxReqLen)
	if err != nil {
		fmt.Printf("Failed to read from connection for edgeDHpublic, err: %s\n", err)
		return
	}
	fmt.Printf("edgeDHPublic: %s\n", edgeDHPublic)
	presharedKey, err := wireguardcommon.GenPskOnManager(edgeDHPublic[:responseLen], managerPrivate)
	if err != nil {
		fmt.Printf("Failed to generate manager private/public key pair: %s\n", err)
		return
	} else {
		fmt.Printf("presharedkey=%s\n", presharedKey)
	}

	///////////////////////////////
	// derive edge keys         //
	/////////////////////////////
	_, edgePublicStr, err := wireguardcommon.CreateKeyPairFromSeed(edgeResponse)
	if err != nil {
		fmt.Printf("Failed to generate manager private/public key pair: %s\n", err)
		return
	}
	///////////////////////////////
	// end of derive edge keys  //
	/////////////////////////////

	var unusedIp netip.Addr
	peer, err := wireguardcommon.FindPeerByPublicKey(edgePublicStr)
	if peer != nil {
		var p = *peer
		unusedIp, err = netip.ParseAddr(p.AllowedIPs[0].IP.String())
	} else {
		unusedIp, err = wireguardcommon.FindUnusedIp()
	}
	allowedIp := unusedIp.String() + "/32"
	if err == nil {
		fmt.Printf("found an unused ip=%s\n", unusedIp)
		secret, _ := base64.StdEncoding.DecodeString(presharedKey)
		ack, _ := wireguardcommon.Chacha20poly1305Encrypt(secret, []byte("OK:"+wireguardcommon.Ip2+":"+allowedIp))

		fmt.Printf("sending ack %x", ack)
		writeLen, err = wireguardcommon.WriteConn(conn, ack)
		if err != nil || writeLen != len(ack) {
			fmt.Printf("write length mismatch. expected = %d, got=%d\n", len(ack), writeLen)
			return
		}

		wireguardcommon.AddPeer(presharedKey, edgePublicStr, "", allowedIp, false)
		/*
			managerIp, _ := wireguardcommon.GetFirstAddr()
			err := wireguardcommon.AddLinkAddress(wireguardLink, managerIp.String(), allowedIp)
			if err != nil {
				fmt.Printf("AddLinkAddress(link=%s, managerIp=%s, allowedIp=%s) failed , error =%s\n", err)
				wireguardcommon.CloseConn(conn)
				return
			}
		*/
	} else {
		fmt.Printf("could not find an unused ip\n")
	}

	err = wireguardcommon.StoreWgConfigUsingWgCommand()
	if err != nil {
		fmt.Printf("Failed to save wg conf, err: %s\n", err)
		os.Exit(wireguardcommon.ExitSetupFailed)
	}

	err = wireguardcommon.StoreWgConfig()
	if err != nil {
		fmt.Printf("Failed to save wg-alt conf, err: %s\n", err)
		os.Exit(wireguardcommon.ExitSetupFailed)
	}

	// Close the connection when you're done with it.
	wireguardcommon.CloseConn(conn)
}
